package org.example.window;

import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.WindowedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingProcessingTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.example.data.WaterSensor;
import org.example.function.WaterSensorMapFunction;

public class WindowAggregateAndProcessDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        //env.enableCheckpointing(2000);

        DataStreamSource<String> source = env.socketTextStream("localhost", 9999);
        SingleOutputStreamOperator<WaterSensor> wsSource = source.map(new WaterSensorMapFunction());

        KeyedStream<WaterSensor, String> sensorKS = wsSource.keyBy(value -> value.getId());

        WindowedStream<WaterSensor, String, TimeWindow> tumblingWindow =
                sensorKS.window(TumblingProcessingTimeWindows.of(Time.seconds(10)));


        SingleOutputStreamOperator<String> aggregate = tumblingWindow.aggregate(new MyAggregateFunction(), new MyProcessFunction());

        aggregate.print();

        env.execute("Window API DEMO");
    }

    public static class MyAggregateFunction implements AggregateFunction<WaterSensor, Integer, String> {
        @Override
        public Integer createAccumulator() {
            System.out.println("创建累加器");
            return 0;
        }

        @Override
        public Integer add(WaterSensor value, Integer accumulator) {
            System.out.println("调用Add方法累加,value:" + (accumulator + value.getVc()));
            return accumulator + value.getVc();
        }

        @Override
        public String getResult(Integer accumulator) {
            System.out.println("调用getResult方法,result:" + accumulator.toString());
            return accumulator.toString();
        }

        @Override
        public Integer merge(Integer a, Integer b) {
            //只有会话窗口会用到，其他窗口不会调用
            System.out.println("调用merge方法");
            return null;
        }
    }

    public static class MyProcessFunction extends ProcessWindowFunction<String, String, String, TimeWindow> {

        @Override
        public void process(String s, Context context, Iterable<String> elements, Collector<String> out) throws Exception {
            long start = context.window().getStart();
            String startDate = DateFormatUtils.format(start, "yyyy-MM-dd HH:mm:ss.SSS");
            long end = context.window().getEnd();
            String endDate = DateFormatUtils.format(end, "yyyy-MM-dd HH:mm:ss.SSS");
            long l = elements.spliterator().estimateSize();
            out.collect("key为：" + s + "的窗口[" + startDate + "," + endDate + "]中一共有" + l + "条数据,elements:" + elements.toString());
        }
    }
}
